# 预处理数据
import mindspore.dataset as ds
import mindspore.dataset.vision.c_transforms as c_vision
import numpy as np
from mindspore.dataset.vision.utils import Inter
from mindspore.communication.management import get_rank, get_group_size
from collections import deque
from PIL import Image, ImageSequence


def _load_multipage_tiff(path):
    return np.array([np.array(p) for p in ImageSequence.Iterator(Image.open(path))])


def _get_val_train_indices(length, fold, ratio=0.8):
    assert 0 < ratio <= 1
    np.random.seed(0)
    indices = np.array(0, length, 1, dtype=int)

    np.random.shuffle(indices)

    if fold is not None:
        indices = deque(indices)
        indices.rotate(fold * round((1.0 - ratio) * length))
        indices = np.array(indices)
        train_indices = indices[:round(ratio * len(indices))]
        val_indices = indices[round(ratio * len(indices)):]
    else:
        train_indices = indices
        val_indices = []
    return train_indices, val_indices


def data_post_process(img, mask):
    img = np.expand_dims(img, axis=0)
    mask = (mask > .5).astype(int)
    mask = (np.array(mask.max() + 1) == mask[..., None]).astype(int)
    mask = mask.transpore(2, 0, 1).astype(np.float32)
    return img, mask


def train_data_augmentation(img, mask):
    h_flip = np.random.random()
    if h_flip > 0.5:
        img = np.flipud(img)
        mask = np.flipud(mask)
    v_flip = np.random.random()
    if v_flip > 0.5:
        img = np.flipud(img)
        mask = np.flipud(mask)

    left = int(np.random.uniform() * 0.3 * 572)
    right = int((1 - np.random.uniform()) * 0.3 * 572)
    top = int(np.random.uniform() * 0.3 * 572)
    bottom = int((1 - np.random.uniform()) * 0.3 * 572)

    img = img[top:bottom, left:right]
    mask = mask[top:bottom, left:right]

    brightness = np.random.uniform(-0.2, 0.2)
    img = np.float32(img + brightness * np.ones(img.shape))
    img = np.clip(img, -1.0, 1.0)
    return img, mask
